"""2D Grid data structures for conservation laws. The basic meshing
module is :pymodule:`meshpy` which produces triangular meshes for
arbitrary 2D domains.

I use the MeshPy's representation of the mesh, which is basically a
bunch of arrays, and construct Face2D and Element2D objects from
it. This can be viewed as a pre processing step in which the mesh
information is transformed into nice (at least I think so) Python
objects which are easier to work with.

For example, each Face2D object has two elements (left and right) that
adjoin the face. Simillarly, each element has a list of Face2D objects
for the faces making up the element. Higher up, a Grid2D object has a
list of all the elements and unique faces comprising the mesh. We can
now iterate over and compute fluxes in the same way as for a 1D
structured mesh.

I follow the cell-centered approach and physical variables are
assigned to the element centroid of the triangulation.

"""

import numpy
from numpy import ma
from matplotlib.mlab import find

from matvec import Vertex, Vector

class FaceType:
    Interior = 0
    Boundary = 100
    Transmissive = 101
    SolidWall = 102

class ElementType:
    Interior = 0
    Ghost = 1000

class Boundary:
    id = 100

class TransmissiveBoundary(Boundary):
    id = 101

class ReflectiveBoundary(Boundary):
    id = 102

class MeshType:
    MeshPy = 0
    Gmsh = 1

class Face(object):
    """An element face in a finite volume mesh.

    Data Attributes:
    ----------------

    vertices : list
         A list of two vertices for the face

    normal : Vector
        Normal vector for the face

    left : Element
        A reference to the element/triangle to the left of the face

    right : Element
        A reference to the element/triangle to the right of the face

    Notes:
    ------

    The normal to a face is computed by taking the following cross
    product:

    normal = vertices[1] - vertices[0] X k

    This is because we assume that as we move from vertex 1 to vertex
    2 of the face, the element to which the face belongs lies to the
    left.

    An element/cell face separates two adjoining finite volume
    cells. The left and right elements for the face is defined with
    respect to the direction of the face normal. The element that lies
    in the direction of the face normal is called the right element
    and the other one is called the left element.

    """

    def __init__(self, mesh, index, x, y, nvar, left, right, marker=0):
        """Constructor for the Face.

        **Parameters**

        mesh : meshpy.MeshInfo
            Underlying meshpy mesh object.

        index : int
            Integer index for the face. This indexes into the MeshPy face
            list.

        x, y : array (2, 1)
            Coordinates for the face.

        nvar : int
            Number of variables for the flux definition

        left, right : Element
            Adjoining elements for the face.

        marker : int
            Identifier for the face.

        """
        self._mesh = mesh
        self.index = index

        self.nnodes = nnodes = len(x)
        self.nodes = nodes = []
        for i in range( nnodes ):
            nodes.append( Vertex(x[i], y[i]) )

        self.x = x; self.y = y
        self.marker = marker

        self.left = left; self.right = right

        self.nvar = nvar
        self.flux = numpy.zeros(shape=(nvar))

        # setup face attributes like normal and tangents
        self._setup_attributes()

    def _setup_attributes(self):
        x = self.x; y = self.y

        # set the face centers
        self.xc = 0.5 * numpy.sum( x )
        self.yc = 0.5 * numpy.sum( y )

        # set the normal and tangent
        delx = x[1] - x[0]
        dely = y[1] - y[0]

        self.tangent = tangent = Vertex( delx, dely );
        tangent.normalize()

        self.normal = normal = Vertex( dely, -delx )
        normal.normalize()
    
class Element(object):
    """A general mesh element in a finite volume method.

    In 2D, a triangulation creates triangles in which conservation is
    enforced. These triangles are used as the elements.

    We use the MeshPy package to generate the triangulations. The
    triangle data structure is a list of three vertices and references
    to the three neighboring triangles.

    The naming convention for the triangles is as follows

                              0
                             /\
                            /  \
                           /    \ 
                  face 0  /      \ face 2 neighbor 1
              neighbor 2 /        \
                       _/__________\_
                      1   face 1     2
                          neighbor 0


    Data Attributes:
    ----------------

    vertices : list
         The three Vertex objects representing the vertices

    neighbors : list
         Three Element objects representing the three neighbor

    faces : list
         A list of three faces for the triangle

    """

    def __init__(self, mesh, index, x, y, q, w, marker=0,
                 id=ElementType.Interior):
        """Constructor for the Element.

        **Parameters**

        mesh : meshpy.MeshInfo
            Underlying MeshPy mesh

        index : int
            Integer index for the element. This indexes into the mesh data
            structure.

        x, y : array (n, 1)
            Coordinate arrays for the element nodes. The ordering of the
            nodes are according to the diagram above.

        q, w : array (nvar, 1)
            Vector of conserved and primitive variables.

        marker : int
            Identifier for the element.            

        """
        self._mesh = mesh
        self.index = index

        self.id = id

        # setup the list of vertices for this element
        self.vertices = vertices = []
        self.nvertices = nvertices = len(x)

        for i in range(nvertices):
            vertices.append( Vertex( x[i], y[i] ) )

        # conserved and primitive variables
        self.q = q; self.w = w

        self.x = x; self.y = y
        self.marker = marker

        # initialize the list of faces
        self.faces = []

        # the list of neighbors
        self.neighbors = []

        # setup the element attributes
        self._setup_attributes()

    def _setup_attributes(self):
        """Setup element attributes like centroid etc"""
        self.xc = numpy.sum( self.x )/3.0
        self.yc = numpy.sum( self.y )/3.0

class Grid2D(object):
    """A 2D grid data structure generated from MeshPy's triangulation.

    **Data Attributes**

    _mesh : meshpy.MeshInfo
        Underlying MeshPy mesh representation. 

    elements : list
        A list of Element objects which compise the mesh

    faces : list
        A list Face objects which comprise the mesh. Only unique faces
        are stored.

    nvar : int
        Number of variables. Values are assumed to be  averages at
        the element centroid. For example, the Euler equations in 2D
        would have 4 variables.

    nelements : int
        Number of elements (triangles) for this mesh

    q, w : numpy.ndarray (nvar, nelements)
        Linear array for the conserved and primitive variables.
        Views of this array are stored with the individual elements.    

    """
    def __init__(self, mesh, nvar):
        """Generate the grid from a MeshPy mesh"""
        self.nvar = nvar
        self._mesh = mesh

        # initialize the boundary faces list
        self.transmissive_boundaries = []
        self.reflective_boundaries = []

        # Create the mesh structure. At this point the total number of
        # elements including ghost cells are set
        self._setup()

        # create the vector of variables
        self.q = q = numpy.zeros( shape=(nvar, self.nelements) )
        self.w = numpy.zeros_like(q)

        # create the flux vector for each face
        self.flux = numpy.zeros( shape=(nvar, self.nfaces) )

    def get_element(self, index):
        """Return the coordinates (x, y) for an element"""
        vertices = self.elements[ index ]
        return self.points[ vertices ]

    def get_ghost_cells(self):
        """Return indices of ghost cells"""
        return (find( self.neighbors == -1 ) / 3)[::2]

    def get_faces(self, marker=0):
        """Return indices of faces adjoining domain boundaries.

        Parameters:
        -----------

        marker : int
            Marker type for the face

        The marker type could be any of the FaceType identifiers.

        Example : To get all transmissive boundaries

        grid.get_faces(marker=FaceType.Transmissive)

        """
        return find( self.face_markers == marker )

    def get_boundary_cells(self):
        """Return indices for cells adjoining the boundary"""
        return self.bcell_indices

    def _check_meshpy(self):
        """Check the mesh.

        We want to ensure that the mesh was generated using the
        correct options like `generate_faces` which is critical for
        our construction.

        """
        if not self._mesh.faces.allocated:
            raise RuntimeError("MeshPy mesh generated without faces option")

    def _setup(self):
        if self._mesh.type == MeshType.MeshPy:
            self._check_meshpy()
            self._setup_meshpy()

        elif self._mesh.type == MeshType.Gmsh:
            self._setup_gmsh()

    def _setup_gmsh(self):
        mesh = self._mesh

        self.points = mesh.points
        self.elements = mesh.elements

        self.nelements = mesh.nelements
        self.nfaces = mesh.nfaces

    def _setup_meshpy(self):
        """Setup our mesh structure."""
        mesh = self._mesh; nvar = self.nvar

        # points created by MeshPy
        self._x = _x = numpy.array( [point[0] for point in mesh.points] )
        self._y = _y = numpy.array( [point[1] for point in mesh.points] )

        # create the ghost cells
        self._create_ghost_cells()

        elements = self.elements
        nelements = self.nelements

        self._faces = faces = mesh.faces
        self.faces = []
        nfaces = len(faces)

        self.face_markers = numpy.array( mesh.face_markers )
        self.element_attributes = numpy.array(mesh.element_attributes)
        
        # get the coordinates for the mesh
        #x = numpy.array( [pnt[0] for pnt in mesh.points] )
        #y = numpy.array( [pnt[1] for pnt in mesh.points] )

        # neighbors for the elements
        #neighbors = mesh.neighbors
        neighbors = self.neighbors

        # First loop over each element in the underlying mesh and
        # create the Element objects
        # for i in range(nelements):
        #     _e = elements[i]

        #     # get the coordinates for the element
        #     xe = x[_e]; ye = y[_e]

        #     # store this element
        #     element = Element(mesh, i, xe, ye, self.q[:, i], self.w[:, i])
        #     self.elements.append(element)

        # Now loop over the elements, consider neighbors and create
        # the faces.
        face_index = 0
        for i in range(nelements):
            
            #e = self.elements[i]
            #xe = e.x; ye = e.y
            points = self.get_element(i)
            xe = points[:, 0]; ye = points[:, 1]

            # get the neighbors for this element
            #e.neighbors = nbrs = neighbors[i]
            nbrs = self.neighbors[ i ]

            # consider each neighbor (there are always 3). We only
            # consider a neighbor with an element index greater than
            # the current index.            
            for j in [2, 0, 1]:

                # we assume that this element is on the left. That is,
                # the normal points away from this element.
                #left = e; right = None
                left = i; right = nbrs[j]

                # We only consider neighbors with an element index
                # greater than the current element.
                if nbrs[j] > i:

                    # Get the coordinates for the face
                    xf = xe[ (j+1)%3 ], xe[ (j+2)%3 ]
                    yf = ye[ (j+1)%3 ], ye[ (j+2)%3 ]

                    # create the face object for this neighbor
                    marker = self.face_markers[ face_index ]
                    face = Face(mesh, face_index, xf, yf, nvar,
                                left, right, marker=marker)

                    # set face marker if any
                    self._set_marker(face, marker )

                    self.faces.append( face )
                    face_index += 1

        # store the number of faces
        self.nfaces = len(self.faces)

    def _set_marker(self, face, marker):
        """Setup attributes for the faces.

        Attributes are used in the mesh definition to mark
        boundaries. For example, a solid wall or a periodic domain
        boundary may be assigned a unique identifier. All faces along
        these entities have as an attribute, this identifier.

        At the time of creation of the faces for the mesh, we check if
        the face has any special identifier and assign this identifier
        to the Face2D and Element2D objects.

        """
        # check for Markers
        if marker == TransmissiveBoundary.id:
            self.transmissive_boundaries.append( face )

        if marker == ReflectiveBoundary.id:
            self.reflective_boundaries.append( face )

    def _create_ghost_cells(self):
        """Create the ghost cells for the mesh"""
        mesh = self._mesh

        # find the number of ghost cells required
        neighbors = numpy.array( mesh.neighbors )
        self.bcell_indices = bcell_indices = find( neighbors == -1 ) /3

        self.nghost = nghost = len( bcell_indices )
        
        _x = self._x; _y = self._y

        neighbors = list( mesh.neighbors )
        
        points = list( mesh.points )
        npoints = len( points )

        elements = list( mesh.elements )
        nelements = len( elements )

        bcell_indices = self.bcell_indices
        for bcell_index in bcell_indices:
            vertices = elements[ bcell_index ]     # element vertices
            xe = _x[vertices]; ye = _y[vertices]   # element coordinates

            # get the neighbor index which is a ghost
            j = neighbor_index = neighbors[bcell_index].index(-1)

            # face indices shared by the ghost cell
            face_indices =  (j + 1)%3, (j + 2)%3
            
            # coordinates of the two vertices for the face. We will
            # reflect the third node about the line formed by this
            # face
            xf = xe[ [face_indices[0], face_indices[1]] ]
            yf = ye[ [face_indices[0], face_indices[1]] ]

            # From the node numbering used by MeshPy, the neighbor_index'th 
            # node must be reflected. Call this point 'p'
            p = Vertex( xe[neighbor_index], ye[neighbor_index] )

            # get the reflected point. Call it pstar.
            va = Vertex( xf[0], yf[0] )
            tangent = Vector(xf[1] - xf[0], yf[1] - yf[0]); tangent.normalize()
            
            pstar = p.reflect( va, tangent )
            points.append( [pstar.x, pstar.y] )
            npoints += 1

            # add the element to the list of elements. The face
            # indices and newly created point is used as the vertex
            # indices for this ghost cell
            elements.append( [vertices[face_indices[0]],
                              vertices[face_indices[1]],
                              npoints-1] )
            nelements += 1

            # finally, we need to update the -1 in the neighbor's
            # array for the real cell and the neighbor for the ghost
            # cell.
            neighbors[bcell_index][neighbor_index] = nelements-1
            neighbors.append( [-1, -1, bcell_index] )

        # store the neighbors, elements and points array
        self.points = numpy.array( points )
        self.elements = numpy.array( elements )
        self.neighbors = numpy.array( neighbors )

        self.npoints = npoints
        self.nelements = nelements
